# -*- coding: UTF-8 -*-
################################################################################
#
# Copyright (c) 2018 Baidu.com, Inc. All Rights Reserved
#
################################################################################
"""
本文件实现了一个完整的demo
"""

from __future__ import absolute_import
from __future__ import print_function
from __future__ import unicode_literals

import multiprocessing
import sys

import cv2
import BaiduAI.EasyEdge as edge
import logging

edge.Log.set_level(logging.INFO)

try:
    _model_dir = sys.argv[1]
    _test_img = sys.argv[2]
except Exception as e:
    print("Usage: python3 demo.py {model_dir} {image_name}")
    exit(-1)
_test_round = 1


edge.set_auth_license_key("936F-DA25-B0CD-DC1D")


def test():
    """单进程"""
    line_color = (125, 175, 75)

    pred = edge.Program()
    
    pred.init(model_dir=_model_dir, device=edge.Device.GPU, engine=edge.Engine.TENSORRT)

    img = cv2.imread(_test_img)
    for i in range(_test_round):
        res = pred.infer_image(img, 0.3)
        h, w, _ = img.shape
        for r in res:
            if pred.model_type == edge.c.ModelType.ObjectDetection:
                x1 = int(r['x1'] * w)
                y1 = int(r['y1'] * h)
                x2 = int(r['x2'] * w)
                y2 = int(r['y2'] * h)
                cv2.rectangle(img, (x1, y1), (x2, y2), line_color, 2)
                cv2.putText(
                    img=img, text=r['label'], org=(x1, y1),
                    fontFace=cv2.FONT_HERSHEY_PLAIN, fontScale=1, color=line_color, thickness=1
                )
            print(r)
        # cv2.imshow("Edge", img)
        cv2.imwrite(_test_img + ".result.jpg", img)
    pred.close()


if __name__ == '__main__':
    test()

